#!/usr/bin/env python
"""
空战环境训练脚本
用于训练多智能体Transformer在空战环境中的表现
"""

import sys
import os
import setproctitle
import numpy as np
from pathlib import Path
import torch

# 添加项目根目录到Python路径
sys.path.append("../../")
from config import get_config  # 导入配置函数
from envs.air_combat.air_combat_env import AirCombatEnv  # 导入空战环境
from envs.env_wrappers import ShareDummyVecEnv  # 导入环境包装器
from runner.shared.air_combat_runner import AIRCOMBATRunner as Runner  # 导入训练运行器


def make_train_env(all_args):
    def get_env_fn(rank):
        def init_env():
            env = AirCombatEnv(all_args)  # 创建空战环境实例
            env.seed(all_args.seed + rank * 1000)  # 设置环境随机种子，确保不同环境有不同的随机性
            return env

        return init_env

    return ShareDummyVecEnv([get_env_fn(0)])


def make_eval_env(all_args):
    def get_env_fn(rank):
        def init_env():
            env = AirCombatEnv(all_args)  # 创建空战环境实例
            env.seed(all_args.seed * 50000 + rank * 10000)  # 使用不同的种子范围，确保评估环境与训练环境不同
            return env

        return init_env

    return ShareDummyVecEnv([get_env_fn(0)])


def parse_args(args, parser):
    # 基本环境参数
    parser.add_argument('--n_agents', type=int, default=4, help="己方单位的数量")  # 己方单位的数量
    parser.add_argument('--n_enemies', type=int, default=4, help="敌方单位的数量")  # 敌方单位的数量
    parser.add_argument('--episode_limit', type=int, default=180, help="每个episode的最大步数")  # 每个episode的最大步数
    parser.add_argument('--run_dir', type=str, default='', help="运行目录")  # 运行目录

    # 解析参数
    all_args = parser.parse_known_args(args)[0]

    return all_args


def main(args):
    # 获取基础配置
    parser = get_config()
    all_args = parse_args(args, parser)

    # 设置计算设备（GPU或CPU）
    if all_args.cuda and torch.cuda.is_available():
        print("choose to use gpu...")  # 选择使用GPU
        device = torch.device("cuda:0")  # 设置GPU设备
        torch.set_num_threads(all_args.n_training_threads)  # 设置PyTorch线程数
        if all_args.cuda_deterministic:
            # 设置CUDA确定性模式，确保结果可重现
            torch.backends.cudnn.benchmark = False
            torch.backends.cudnn.deterministic = True
    else:
        print("choose to use cpu...")  # 选择使用CPU
        device = torch.device("cpu")
        torch.set_num_threads(all_args.n_training_threads)

    # 创建结果保存目录
    # 构建目录路径：results/
    run_dir = Path(os.path.split(os.path.dirname(os.path.abspath(__file__)))[
                       0] + "/results")
    if not run_dir.exists():
        os.makedirs(str(run_dir))  # 创建目录

    # 确定当前运行的编号
    if not run_dir.exists():
        curr_run = 'run1'  # 第一次运行
    else:
        # 查找已存在的运行编号
        exst_run_nums = [int(str(folder.name).split('run')[1]) for folder in run_dir.iterdir() if
                         str(folder.name).startswith('run')]
        if len(exst_run_nums) == 0:
            curr_run = 'run1'  # 没有找到现有运行
        else:
            curr_run = 'run%i' % (max(exst_run_nums) + 1)  # 使用下一个编号
    run_dir = run_dir / curr_run  # 添加运行编号到路径
    if not run_dir.exists():
        os.makedirs(str(run_dir))  # 创建最终运行目录

    # 设置进程标题，便于在系统监控中识别
    setproctitle.setproctitle("MAT")

    # 设置随机种子，确保结果可重现
    torch.manual_seed(all_args.seed)  # PyTorch CPU随机种子
    torch.cuda.manual_seed_all(all_args.seed)  # PyTorch GPU随机种子
    np.random.seed(all_args.seed)  # NumPy随机种子

    # 创建环境
    num_agents = all_args.n_agents  # 获取智能体数量（4架无人机）
    all_args.run_dir = run_dir  # 将运行目录添加到参数中
    envs = make_train_env(all_args)  # 创建训练环境
    eval_envs = make_eval_env(all_args) if all_args.use_eval else None  # 创建评估环境（如果启用）

    # 构建配置字典
    config = {
        "all_args": all_args,  # 所有参数
        "envs": envs,  # 训练环境
        "eval_envs": eval_envs,  # 评估环境
        "num_agents": num_agents,  # 智能体数量
        "device": device,  # 计算设备
        "run_dir": run_dir  # 运行目录
    }

    # 创建训练运行器并开始训练
    runner = Runner(config)
    runner.run()

    # 训练后的清理工作
    envs.close()  # 关闭训练环境
    if all_args.use_eval and eval_envs is not envs:
        eval_envs.close()  # 关闭评估环境（如果不是同一个环境）

    # 导出训练日志
    runner.writter.export_scalars_to_json(str(runner.log_dir + '/summary.json'))  # 导出标量数据到JSON
    runner.writter.close()  # 关闭日志写入器


if __name__ == "__main__":
    main(sys.argv[1:])  # 传入除脚本名外的所有命令行参数
